#include "visitor/BasicVisitor.h"
#include "AST/UniASTNode.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "symbolTable/Expression.h"
#include "symbolTable/SymbolTable.h"
#include "symbolTable/Variable.h"
#include "utils/Logger.h"
#include <any>
#include <memory>

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Verifier.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/ScopedHashTable.h"
#include "llvm/ADT/StringRef.h"

#define REGISTER_FUNCTION(name, visitFunction)                            \
    this->registerFunction(name, [this](TreeNode *node)                   \
                           { return visitFunction(node); })

UniAST::BasicVisitor::BasicVisitor()
    : manager(Manager::getManager()),
      builder(mlir::OpBuilder(&this->mlirContext)),
      theModule(mlir::ModuleOp::create(builder.getUnknownLoc()))
{
    LOGGER.setLogLevel(Logger::LogLevel::DEBUG);
    // register functions
    REGISTER_FUNCTION("root", this->visitRoot);
    REGISTER_FUNCTION("module", this->visitModule);
    REGISTER_FUNCTION("assign_stmt", this->visitAssign_stmt);
    REGISTER_FUNCTION("variable_access", this->visitVariable_access);
    REGISTER_FUNCTION("name_access", this->visitName_access);
    REGISTER_FUNCTION("expression", this->visitExpression);
    REGISTER_FUNCTION("op", this->visitOp);
    REGISTER_FUNCTION("expr_params", this->visitExpr_params);
    REGISTER_FUNCTION("expr_param", this->visitExpr_param);
    REGISTER_FUNCTION("NUMBER", this->visitNUMBER);
    REGISTER_FUNCTION("TEXT", this->visitTEXT);
    REGISTER_FUNCTION("literal", this->visitLiteral);
    REGISTER_FUNCTION("ADDOP", this->visitADDOP);
    REGISTER_FUNCTION("SUBOP", this->visitSUBOP);

    // initialize mlir
    this->mlirContext.getOrLoadDialect<mlir::arith::ArithDialect>();
    this->mlirContext.getOrLoadDialect<mlir::memref::MemRefDialect>();
    this->mlirContext.getOrLoadDialect<mlir::func::FuncDialect>();
    this->builder.setInsertionPointToEnd(this->theModule.getBody());
    auto loc = this->builder.getUnknownLoc();
    auto mainfunc = this->builder.create<mlir::func::FuncOp>(
        loc, "main",
        this->builder.getFunctionType({}, {builder.getI32Type()}));
    mainfunc.setPrivate();
    auto mainbody = mainfunc.addEntryBlock();
    this->builder.setInsertionPointToEnd(mainbody);
}

#undef REGISTER_FUNCTION

UniAST::BasicVisitor::~BasicVisitor()
{
    auto loc = this->builder.getUnknownLoc();
    auto c0 = this->builder.create<mlir::arith::ConstantIntOp>(loc, 0, 32);
    this->builder.create<mlir::func::ReturnOp>(loc, mlir::ValueRange{c0});
    if (llvm::failed(mlir::verify(this->theModule)))
    {
        LOGGER.log(Logger::LogLevel::ERROR, "verify failed.");
    }
    this->theModule->dump();
}

void UniAST::BasicVisitor::registerFunction(
    std::string nodeName,
    std::function<std::any(UniAST::TreeNode *)> visitFunction)
{
    this->visitFunctions[nodeName] = visitFunction;
}

std::any UniAST::BasicVisitor::visit(TreeNode *node)
{
    if (!node)
    {
        LOGGER.log(Logger::LogLevel::ERROR, "empty node!");
        exit(1);
    }
    auto nodeName = node->getNodeName();
    if (this->visitFunctions.contains(nodeName))
    {
        return this->visitFunctions[nodeName](node);
    }
    else
    {
        LOGGER.log(Logger::LogLevel::WARN,
                   "Can not find visitor for node: ")
            << nodeName << ", use default visit function. ";
        if (node->getCtx() != nullptr)
        {
            LOGGER << "Context details: "
                   << node->getCtx()->toStringTree();
        }
        LOGGER << "\n";
        return this->visitChildren(node);
    }
}

std::any UniAST::BasicVisitor::visitChildren(TreeNode *node)
{
    std::any result;
    for (auto child : node->getChildren())
    {
        result = this->visit(child);
    }
    return result;
}

std::any UniAST::BasicVisitor::visitRoot(TreeNode *node)
{
    return this->visitChildren(node);
}

std::any UniAST::BasicVisitor::visitModule(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit module. \n");
    return this->visitChildren(node);
}

mlir::Location UniAST::BasicVisitor::getLocation(UniAST::TreeNode *node)
{
    auto ctx = node->getCtx();
    auto line = ctx->start->getLine();
    auto col = ctx->start->getCharPositionInLine();
    return mlir::FileLineColLoc::get(
        this->builder.getStringAttr("source_file"), line, col);
}

std::any UniAST::BasicVisitor::visitAssign_stmt(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Assign_stmt. \n");
    auto loc = this->getLocation(node);
    auto var = std::any_cast<std::shared_ptr<Variable>>(
        this->visit(node->findChild("variable_access")));
    var->defined = true;

    auto expr = std::any_cast<std::shared_ptr<RightValue>>(
        this->visit(node->findChild("expression")));

    // set variable type
    auto varty = expr->getType();
    var->setType(varty);

    // create alloca op
    auto elemty = varty->getmlirType(this->builder);
    auto memrefty =
        mlir::MemRefType::get(llvm::ArrayRef<int64_t>{1}, elemty);
    auto alloca =
        this->builder.create<mlir::memref::AllocaOp>(loc, memrefty);
    var->setValue(alloca);
    // get expr mlir value
    auto exprValue = expr->getmlirValue();
    // create store
    auto c0 = this->builder.create<mlir::arith::ConstantIndexOp>(loc, 0);
    this->builder.create<mlir::memref::StoreOp>(loc, exprValue, alloca,
                                                mlir::ValueRange{c0});

    // variable do not exist
    // add variable to symbol table
    auto varname = var->getName();
    if (!this->manager.currentScope()->existVariable(varname))
    {
        this->manager.currentScope()->addVariable(var);
    }
    return std::any();
}

std::any UniAST::BasicVisitor::visitVariable_access(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Variable_access. \n");
    return this->visitChildren(node);
}

std::any UniAST::BasicVisitor::visitName_access(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Name_access. \n");
    auto varname =
        std::any_cast<std::string>(this->visit(node->findChild("TEXT")));
    if (this->manager.currentScope()->existVariable(varname))
    {
        return this->manager.currentScope()->findVariableFromThis(varname);
    }
    // create new variable
    return Variable::get(varname);
}

std::any UniAST::BasicVisitor::visitExpression(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Expression. \n");
    auto loc = this->getLocation(node);
    auto expr_params =
        std::any_cast<std::vector<std::shared_ptr<RightValue>>>(
            this->visit(node->findChild("expr_params")));
    auto op =
        std::any_cast<std::string>(this->visit(node->findChild("op")));
    if (op == "+")
    {
        auto res = Expression::get(expr_params, Expression::OPERATOR::ADD);
        auto resTy = res->getType();
        auto lhs = expr_params[0];
        auto rhs = expr_params[1];
        auto lhsmlirvalue = lhs->getmlirValue();
        auto rhsmlirvalue = rhs->getmlirValue();
        mlir::Value resmlirvalue;
        if (resTy->isFloatType())
        {
            resmlirvalue = this->builder.create<mlir::arith::AddFOp>(
                loc, lhsmlirvalue, rhsmlirvalue);
        }
        else if (resTy->isIntegerType())
        {
            resmlirvalue = this->builder.create<mlir::arith::AddIOp>(
                loc, lhsmlirvalue, rhsmlirvalue);
        }
        res->setmlirValue(resmlirvalue);
        return std::static_pointer_cast<RightValue>(res);
    }
    else if (op == "-")
    {
        auto res = Expression::get(expr_params, Expression::OPERATOR::SUB);
        auto resTy = res->getType();
        auto lhs = expr_params[0];
        auto rhs = expr_params[1];
        auto lhsmlirvalue = lhs->getmlirValue();
        auto rhsmlirvalue = rhs->getmlirValue();
        mlir::Value resmlirvalue;
        if (resTy->isFloatType())
        {
            resmlirvalue = this->builder.create<mlir::arith::SubFOp>(
                loc, lhsmlirvalue, rhsmlirvalue);
        }
        else if (resTy->isIntegerType())
        {
            resmlirvalue = this->builder.create<mlir::arith::SubIOp>(
                loc, lhsmlirvalue, rhsmlirvalue);
        }
        res->setmlirValue(resmlirvalue);
        return std::static_pointer_cast<RightValue>(res);
    }
    else
    {
        assert(false && "unknown op");
        exit(1);
    }
}

std::any UniAST::BasicVisitor::visitOp(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Op. \n");
    return this->visitChildren(node);
}

std::any UniAST::BasicVisitor::visitExpr_params(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Expr_params. \n");
    std::vector<std::shared_ptr<RightValue>> res;
    // collect all children and returns
    for (auto child : node->findChildren("expr_param"))
    {
        res.push_back(std::any_cast<std::shared_ptr<RightValue>>(
            this->visit(child)));
    }
    return res;
}

std::any UniAST::BasicVisitor::visitExpr_param(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit Expr_param. \n");
    auto loc = this->getLocation(node);
    auto nodeName = node->getChildren()[0]->getNodeName();
    auto childret = this->visit(node->getChildren()[0]);
    if (nodeName == "expression")
    {
        return childret;
    }
    else if (nodeName == "variable_access")
    {
        auto var = std::any_cast<std::shared_ptr<Variable>>(childret);
        auto varName = var->getName();
        if (!var->defined)
        {
            LOGGER.log(Logger::LogLevel::ERROR, "undefined variable: ")
                << varName;
            exit(1);
        }
        auto c1 = builder.create<mlir::arith::ConstantIndexOp>(loc, 1);
        auto load = builder.create<mlir::memref::LoadOp>(
            loc, var->getValue(), mlir::ValueRange{c1});
        return std::static_pointer_cast<RightValue>(
            VariableAccess::get(var, load));
    }
    else
    {
        auto literal =
            std::any_cast<std::shared_ptr<UniAST::Literal>>(childret);
        return std::static_pointer_cast<RightValue>(literal);
    }
}

std::any UniAST::BasicVisitor::visitLiteral(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit literal. \n");
    auto loc = this->getLocation(node);
    auto number = std::any_cast<int>(this->visit(node->getChildren()[0]));
    auto mlirValue =
        this->builder.create<mlir::arith::ConstantIntOp>(loc, number, 32);
    auto intliteral = UniAST::IntergerLiteral::get(number, mlirValue);
    return std::static_pointer_cast<UniAST::Literal>(intliteral);
}

std::any UniAST::BasicVisitor::visitNUMBER(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit NUMBER. \n");
    auto numbertext = node->getCtx()->getText();
    return int(std::stoi(numbertext));
}

std::any UniAST::BasicVisitor::visitTEXT(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit TEXT. \n");
    return node->getCtx()->getText();
}

std::any UniAST::BasicVisitor::visitADDOP(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit add. \n");
    return node->getCtx()->getText();
}

std::any UniAST::BasicVisitor::visitSUBOP(TreeNode *node)
{
    LOGGER.log(Logger::LogLevel::DEBUG, "visit sub. \n");
    return node->getCtx()->getText();
}